"""
模型训练脚本
训练XGBoost和LSTM模型用于慢性病风险预测
"""

import pandas as pd
import numpy as np
import os
import sys
from datetime import datetime

# 添加项目路径
sys.path.append('.')

from models.xgboost_model import XGBoostPredictor
from models.lstm_model import LSTMModelWrapper
from data.preprocessing import DataPreprocessor
from demo_data_generator import DemoDataGenerator

class ModelTrainer:
    """模型训练器"""
    
    def __init__(self):
        self.preprocessor = DataPreprocessor()
        self.xgboost_model = XGBoostPredictor()
        self.lstm_model = LSTMModelWrapper()
        
    def prepare_training_data(self, data_file='demo_data/patient_data.csv'):
        """准备训练数据"""
        print("正在加载训练数据...")
        
        if not os.path.exists(data_file):
            print(f"数据文件不存在: {data_file}")
            print("正在生成演示数据...")
            generator = DemoDataGenerator()
            patient_data, _ = generator.save_demo_data()
        else:
            patient_data = pd.read_csv(data_file)
        
        print(f"加载了 {len(patient_data)} 条患者数据")
        
        # 准备特征和目标变量
        feature_columns = [
            'age', 'gender', 'bmi', 'systolic_bp', 'diastolic_bp',
            'blood_glucose', 'cholesterol', 'heart_rate', 'smoking',
            'drinking', 'family_history', 'exercise_frequency', 'diet_preference'
        ]
        
        X = patient_data[feature_columns]
        y_hypertension = patient_data['hypertension_risk']
        y_diabetes = patient_data['diabetes_risk']
        
        print(f"特征维度: {X.shape}")
        print(f"高血压阳性样本: {y_hypertension.sum()} ({y_hypertension.mean()*100:.1f}%)")
        print(f"糖尿病阳性样本: {y_diabetes.sum()} ({y_diabetes.mean()*100:.1f}%)")
        
        return X, y_hypertension, y_diabetes
    
    def prepare_timeseries_data(self, data_file='demo_data/timeseries_data.csv'):
        """准备时序训练数据"""
        print("正在加载时序数据...")
        
        if not os.path.exists(data_file):
            print(f"时序数据文件不存在: {data_file}")
            return None
        
        timeseries_data = pd.read_csv(data_file)
        print(f"加载了 {len(timeseries_data)} 条时序数据")
        
        # 按患者分组处理时序数据
        patient_groups = timeseries_data.groupby('patient_id')
        
        # 选择数据量足够的患者
        valid_patients = []
        for patient_id, group in patient_groups:
            if len(group) >= 30:  # 至少30天的数据
                valid_patients.append(group)
        
        if len(valid_patients) == 0:
            print("没有足够的时序数据进行训练")
            return None
        
        print(f"找到 {len(valid_patients)} 个患者有足够的时序数据")
        
        # 合并所有患者的时序数据
        combined_data = pd.concat(valid_patients, ignore_index=True)
        
        return combined_data
    
    def train_xgboost_models(self, X, y_hypertension, y_diabetes):
        """训练XGBoost模型"""
        print("\n开始训练XGBoost模型...")
        print("-" * 40)
        
        try:
            self.xgboost_model.train(X, y_hypertension, y_diabetes)
            
            # 保存模型
            self.xgboost_model.save_models('models/xgboost_models.pkl')
            
            print("XGBoost模型训练完成并已保存")
            return True
            
        except Exception as e:
            print(f"XGBoost模型训练失败: {e}")
            return False
    
    def train_lstm_model(self, timeseries_data):
        """训练LSTM模型"""
        print("\n开始训练LSTM模型...")
        print("-" * 40)
        
        if timeseries_data is None:
            print("跳过LSTM模型训练（无时序数据）")
            return False
        
        try:
            # 准备时序数据
            timeseries_data['date'] = pd.to_datetime(timeseries_data['date'])
            timeseries_data = timeseries_data.sort_values(['patient_id', 'date'])
            
            # 训练模型
            self.lstm_model.train(timeseries_data, epochs=50, batch_size=16, learning_rate=0.001)
            
            # 保存模型
            self.lstm_model.save_model('models/lstm_model.pth')
            
            print("LSTM模型训练完成并已保存")
            return True
            
        except Exception as e:
            print(f"LSTM模型训练失败: {e}")
            return False
    
    def evaluate_models(self, X, y_hypertension, y_diabetes):
        """评估模型性能"""
        print("\n模型性能评估...")
        print("-" * 40)
        
        try:
            # 评估XGBoost模型
            if self.xgboost_model.is_trained:
                print("XGBoost模型评估:")
                
                # 高血压预测
                ht_pred = self.xgboost_model.hypertension_model.predict(X)
                ht_proba = self.xgboost_model.hypertension_model.predict_proba(X)[:, 1]
                
                from sklearn.metrics import accuracy_score, precision_score, recall_score, f1_score, roc_auc_score
                
                ht_accuracy = accuracy_score(y_hypertension, ht_pred)
                ht_precision = precision_score(y_hypertension, ht_pred)
                ht_recall = recall_score(y_hypertension, ht_pred)
                ht_f1 = f1_score(y_hypertension, ht_pred)
                ht_auc = roc_auc_score(y_hypertension, ht_proba)
                
                print(f"  高血压预测:")
                print(f"    准确率: {ht_accuracy:.3f}")
                print(f"    精确率: {ht_precision:.3f}")
                print(f"    召回率: {ht_recall:.3f}")
                print(f"    F1分数: {ht_f1:.3f}")
                print(f"    AUC: {ht_auc:.3f}")
                
                # 糖尿病预测
                dm_pred = self.xgboost_model.diabetes_model.predict(X)
                dm_proba = self.xgboost_model.diabetes_model.predict_proba(X)[:, 1]
                
                dm_accuracy = accuracy_score(y_diabetes, dm_pred)
                dm_precision = precision_score(y_diabetes, dm_pred)
                dm_recall = recall_score(y_diabetes, dm_pred)
                dm_f1 = f1_score(y_diabetes, dm_pred)
                dm_auc = roc_auc_score(y_diabetes, dm_proba)
                
                print(f"  糖尿病预测:")
                print(f"    准确率: {dm_accuracy:.3f}")
                print(f"    精确率: {dm_precision:.3f}")
                print(f"    召回率: {dm_recall:.3f}")
                print(f"    F1分数: {dm_f1:.3f}")
                print(f"    AUC: {dm_auc:.3f}")
                
                # 特征重要性
                print("\n特征重要性分析:")
                importance = self.xgboost_model.get_feature_importance()
                if importance:
                    print("  高血压预测特征重要性:")
                    for feature, imp in sorted(importance['hypertension'].items(), key=lambda x: x[1], reverse=True)[:5]:
                        print(f"    {feature}: {imp:.3f}")
                    
                    print("  糖尿病预测特征重要性:")
                    for feature, imp in sorted(importance['diabetes'].items(), key=lambda x: x[1], reverse=True)[:5]:
                        print(f"    {feature}: {imp:.3f}")
            
            # 评估LSTM模型
            if self.lstm_model.is_trained:
                print("\nLSTM模型评估:")
                print("  LSTM模型主要用于时序预测，具体评估需要时序测试数据")
            
        except Exception as e:
            print(f"模型评估失败: {e}")
    
    def generate_model_report(self):
        """生成模型训练报告"""
        report_file = 'models/training_report.txt'
        os.makedirs('models', exist_ok=True)
        
        with open(report_file, 'w', encoding='utf-8') as f:
            f.write("慢性病风险预测模型训练报告\n")
            f.write("=" * 50 + "\n\n")
            
            f.write(f"训练时间: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}\n\n")
            
            f.write("模型状态:\n")
            f.write("-" * 30 + "\n")
            f.write(f"XGBoost模型: {'已训练' if self.xgboost_model.is_trained else '未训练'}\n")
            f.write(f"LSTM模型: {'已训练' if self.lstm_model.is_trained else '未训练'}\n\n")
            
            f.write("模型文件:\n")
            f.write("-" * 30 + "\n")
            f.write("- models/xgboost_models.pkl (XGBoost模型)\n")
            f.write("- models/lstm_model.pth (LSTM模型)\n\n")
            
            f.write("使用说明:\n")
            f.write("-" * 30 + "\n")
            f.write("1. 启动Flask应用: python app.py\n")
            f.write("2. 访问系统: http://localhost:5000\n")
            f.write("3. 使用训练好的模型进行风险预测\n")
            f.write("4. 查看SHAP解释和干预建议\n")
        
        print(f"训练报告已保存到: {report_file}")
    
    def train_all_models(self):
        """训练所有模型"""
        print("慢性病风险预测模型训练")
        print("=" * 50)
        
        # 准备数据
        X, y_hypertension, y_diabetes = self.prepare_training_data()
        timeseries_data = self.prepare_timeseries_data()
        
        # 训练XGBoost模型
        xgb_success = self.train_xgboost_models(X, y_hypertension, y_diabetes)
        
        # 训练LSTM模型
        lstm_success = self.train_lstm_model(timeseries_data)
        
        # 评估模型
        if xgb_success:
            self.evaluate_models(X, y_hypertension, y_diabetes)
        
        # 生成报告
        self.generate_model_report()
        
        print("\n模型训练完成！")
        print(f"XGBoost模型: {'成功' if xgb_success else '失败'}")
        print(f"LSTM模型: {'成功' if lstm_success else '失败'}")
        
        if xgb_success or lstm_success:
            print("\n现在可以启动Flask应用进行测试:")
            print("python app.py")

def main():
    """主函数"""
    trainer = ModelTrainer()
    trainer.train_all_models()

if __name__ == "__main__":
    main()



